package reverse_proxy

import (
	"context"
	"io"
	"log"
	"net"
	"time"
)

func NewTcpReverseProxy(ctx context.Context, addr string) *TcpReverseProxy {
	return &TcpReverseProxy{
		ctx:             ctx,
		Addr:            addr,
		KeepAlivePeriod: time.Second,
		DialTimeout:     time.Second,
	}
}

type TcpReverseProxy struct {
	ctx             context.Context
	Addr            string
	KeepAlivePeriod time.Duration
	DailContext     func(ctx context.Context, network, address string) (net.Conn, error)
	DialTimeout     time.Duration
	OnDialError     func(src net.Conn, dstDialErr error)
}

func (tp *TcpReverseProxy) dailTimeout() time.Duration {
	if tp.DialTimeout > 0 {
		return tp.DialTimeout
	}
	return 10 * time.Second
}

func (tp *TcpReverseProxy) dailContext() func(ctx context.Context, network, address string) (net.Conn, error) {
	if tp.DailContext != nil {
		return tp.DailContext
	}
	return (&net.Dialer{
		Timeout:   tp.DialTimeout,
		KeepAlive: tp.KeepAlivePeriod,
	}).DialContext
}

func (tp *TcpReverseProxy) onDialError() func(net.Conn, error) {
	if tp.OnDialError != nil {
		return tp.OnDialError
	}
	return func(src net.Conn, dstDialErr error) {
		log.Printf("tcpproxy: for incoming conn %v, error dialing %q: %v", src.RemoteAddr().String(), tp.Addr, dstDialErr)
		src.Close()
	}
}

func (tp *TcpReverseProxy) keepAlivePeriod() time.Duration {
	if tp.KeepAlivePeriod != 0 {
		return tp.KeepAlivePeriod
	}
	return time.Minute
}

func (tp *TcpReverseProxy) ServeTCP(ctx context.Context, src net.Conn) {
	var cancel context.CancelFunc
	if tp.DialTimeout >= 0 {
		ctx, cancel = context.WithTimeout(ctx, tp.dailTimeout())
	}
	dst, err := tp.dailContext()(ctx, "tcp", tp.Addr)
	if cancel != nil {
		cancel()
	}
	if err != nil {
		tp.onDialError()(src, err)
		return
	}

	if ka := tp.keepAlivePeriod(); ka > 0 {
		if c, ok := dst.(*net.TCPConn); ok {
			c.SetKeepAlive(true)
			c.SetKeepAlivePeriod(ka)
		}
	}

	errCh := make(chan error, 1)
	go tp.proxyCopy(errCh, src, dst)
	go tp.proxyCopy(errCh, dst, src)
	<-errCh

}

func (tp *TcpReverseProxy) proxyCopy(errCh chan<- error, dst, src net.Conn) {
	_, err := io.Copy(dst, src)
	errCh <- err
}
